#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
文件组织模块 - 负责根据不同规则整理文件
"""
import os
import shutil
import datetime
from pathlib import Path
from typing import Dict, List, Union, Optional, Callable

# 常见文件类型映射
FILE_TYPE_EXTENSIONS = {
    "图片": [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"],
    "文档": [".doc", ".docx", ".pdf", ".txt", ".rtf", ".md", ".odt", ".xlsx", ".xls", ".ppt", ".pptx"],
    "视频": [".mp4", ".avi", ".mov", ".wmv", ".flv", ".mkv", ".webm", ".m4v"],
    "音频": [".mp3", ".wav", ".ogg", ".flac", ".aac", ".wma", ".m4a"],
    "压缩包": [".zip", ".rar", ".7z", ".tar", ".gz", ".bz2"],
    "代码": [".py", ".js", ".html", ".css", ".java", ".cpp", ".c", ".php", ".go", ".rb"],
    "可执行文件": [".exe", ".msi", ".app", ".dmg", ".sh", ".bat"],
}

def organize_by_type(source_dir: Union[str, Path], target_dir: Optional[Union[str, Path]] = None, 
                     file_types: Optional[Dict[str, List[str]]] = None) -> Dict[str, int]:
    """
    按文件类型整理文件
    
    Args:
        source_dir: 源文件夹路径
        target_dir: 目标文件夹路径（如果为None，则在源文件夹中创建子文件夹）
        file_types: 自定义文件类型映射（如果为None，则使用默认映射）
    
    Returns:
        Dict[str, int]: 每种类型移动的文件数量
    """
    source_path = Path(source_dir)
    if not source_path.exists() or not source_path.is_dir():
        raise ValueError(f"源目录不存在或不是文件夹: {source_dir}")
    
    # 如果没有指定目标目录，则在源目录中创建
    if target_dir is None:
        target_dir = source_path
    else:
        target_dir = Path(target_dir)
        if not target_dir.exists():
            target_dir.mkdir(parents=True)
    
    # 使用默认文件类型或自定义类型
    types_map = file_types if file_types is not None else FILE_TYPE_EXTENSIONS
    
    # 用于存储每种类型移动的文件数
    stats = {category: 0 for category in types_map}
    stats["其他"] = 0
    
    # 查找源目录中的所有文件（不包括子目录中的文件）
    files = [f for f in source_path.iterdir() if f.is_file()]
    
    # 按类型整理
    for file_path in files:
        file_ext = file_path.suffix.lower()
        target_category = "其他"
        
        # 确定文件类型
        for category, extensions in types_map.items():
            if file_ext in extensions:
                target_category = category
                break
        
        # 创建目标子目录（如果不存在）
        target_subdir = target_dir / target_category
        if not target_subdir.exists():
            target_subdir.mkdir(parents=True)
        
        # 移动文件
        try:
            target_file = target_subdir / file_path.name
            if target_file.exists():
                # 处理重名文件
                name_parts = file_path.stem, file_path.suffix
                target_file = target_subdir / f"{name_parts[0]}_copy{name_parts[1]}"
            
            shutil.move(str(file_path), str(target_file))
            stats[target_category] += 1
        except Exception as e:
            print(f"移动文件 {file_path} 时出错: {e}")
    
    return stats

def organize_by_date(source_dir: Union[str, Path], target_dir: Optional[Union[str, Path]] = None,
                    date_format: str = "年份/月份", use_modified_date: bool = False) -> Dict[str, int]:
    """
    按日期整理文件
    
    Args:
        source_dir: 源文件夹路径
        target_dir: 目标文件夹路径（如果为None，则在源文件夹中创建子文件夹）
        date_format: 日期格式，可选："年份", "年份/月份", "年份/月份/日期"
        use_modified_date: 是否使用修改日期（否则使用创建日期）
    
    Returns:
        Dict[str, int]: 每个时间段移动的文件数量
    """
    source_path = Path(source_dir)
    if not source_path.exists() or not source_path.is_dir():
        raise ValueError(f"源目录不存在或不是文件夹: {source_dir}")
    
    # 如果没有指定目标目录，则在源目录中创建
    if target_dir is None:
        target_dir = source_path
    else:
        target_dir = Path(target_dir)
        if not target_dir.exists():
            target_dir.mkdir(parents=True)
    
    # 用于存储每个时间段移动的文件数
    stats = {}
    
    # 查找源目录中的所有文件（不包括子目录中的文件）
    files = [f for f in source_path.iterdir() if f.is_file()]
    
    # 按日期整理
    for file_path in files:
        # 获取文件的创建时间或修改时间
        if use_modified_date:
            file_time = datetime.datetime.fromtimestamp(os.path.getmtime(file_path))
        else:
            file_time = datetime.datetime.fromtimestamp(os.path.getctime(file_path))
        
        year = str(file_time.year)
        month = f"{file_time.month:02d}"
        day = f"{file_time.day:02d}"
        
        # 根据指定的日期格式创建目标路径
        if date_format == "年份":
            date_path = year
            period = year
        elif date_format == "年份/月份":
            date_path = os.path.join(year, month)
            period = f"{year}-{month}"
        elif date_format == "年份/月份/日期":
            date_path = os.path.join(year, month, day)
            period = f"{year}-{month}-{day}"
        else:
            raise ValueError(f"不支持的日期格式: {date_format}")
        
        target_subdir = target_dir / date_path
        
        # 初始化统计数据
        if period not in stats:
            stats[period] = 0
        
        # 创建目标子目录（如果不存在）
        if not target_subdir.exists():
            target_subdir.mkdir(parents=True)
        
        # 移动文件
        try:
            target_file = target_subdir / file_path.name
            if target_file.exists():
                # 处理重名文件
                name_parts = file_path.stem, file_path.suffix
                target_file = target_subdir / f"{name_parts[0]}_copy{name_parts[1]}"
                
            shutil.move(str(file_path), str(target_file))
            stats[period] += 1
        except Exception as e:
            print(f"移动文件 {file_path} 时出错: {e}")
    
    return stats

def organize_by_size(source_dir: Union[str, Path], target_dir: Optional[Union[str, Path]] = None,
                    size_categories: Optional[Dict[str, int]] = None) -> Dict[str, int]:
    """
    按文件大小整理文件
    
    Args:
        source_dir: 源文件夹路径
        target_dir: 目标文件夹路径（如果为None，则在源文件夹中创建子文件夹）
        size_categories: 自定义大小分类（以字节为单位），例如：{"小": 1024*1024, "中": 10*1024*1024, "大": float("inf")}
    
    Returns:
        Dict[str, int]: 每个大小类别移动的文件数量
    """
    source_path = Path(source_dir)
    if not source_path.exists() or not source_path.is_dir():
        raise ValueError(f"源目录不存在或不是文件夹: {source_dir}")
    
    # 如果没有指定目标目录，则在源目录中创建
    if target_dir is None:
        target_dir = source_path
    else:
        target_dir = Path(target_dir)
        if not target_dir.exists():
            target_dir.mkdir(parents=True)
    
    # 默认大小分类（小：<1MB，中：1MB-100MB，大：>100MB）
    if size_categories is None:
        size_categories = {
            "小文件": 1024 * 1024,  # 1MB
            "中等文件": 100 * 1024 * 1024,  # 100MB
            "大文件": float("inf")
        }
    
    # 排序分类边界，确保从小到大排序
    sorted_categories = sorted(size_categories.items(), key=lambda x: x[1])
    
    # 用于存储每个大小类别移动的文件数
    stats = {category: 0 for category in size_categories}
    
    # 查找源目录中的所有文件（不包括子目录中的文件）
    files = [f for f in source_path.iterdir() if f.is_file()]
    
    # 按大小整理
    for file_path in files:
        file_size = file_path.stat().st_size
        
        # 确定文件大小类别
        target_category = sorted_categories[-1][0]  # 默认为最大的类别
        for category, size_limit in sorted_categories:
            if file_size < size_limit:
                target_category = category
                break
        
        # 创建目标子目录（如果不存在）
        target_subdir = target_dir / target_category
        if not target_subdir.exists():
            target_subdir.mkdir(parents=True)
        
        # 移动文件
        try:
            target_file = target_subdir / file_path.name
            if target_file.exists():
                # 处理重名文件
                name_parts = file_path.stem, file_path.suffix
                target_file = target_subdir / f"{name_parts[0]}_copy{name_parts[1]}"
                
            shutil.move(str(file_path), str(target_file))
            stats[target_category] += 1
        except Exception as e:
            print(f"移动文件 {file_path} 时出错: {e}")
    
    return stats

def organize_by_regex(source_dir: Union[str, Path], target_dir: Optional[Union[str, Path]] = None,
                     patterns: Dict[str, str] = None) -> Dict[str, int]:
    """
    按正则表达式匹配文件名整理文件
    
    Args:
        source_dir: 源文件夹路径
        target_dir: 目标文件夹路径（如果为None，则在源文件夹中创建子文件夹）
        patterns: 正则表达式模式字典，格式为 {"类别名": "正则表达式"}
    
    Returns:
        Dict[str, int]: 每个类别移动的文件数量
    """
    import re
    
    source_path = Path(source_dir)
    if not source_path.exists() or not source_path.is_dir():
        raise ValueError(f"源目录不存在或不是文件夹: {source_dir}")
    
    # 如果没有指定目标目录，则在源目录中创建
    if target_dir is None:
        target_dir = source_path
    else:
        target_dir = Path(target_dir)
        if not target_dir.exists():
            target_dir.mkdir(parents=True)
    
    # 验证正则表达式字典
    if not patterns or not isinstance(patterns, dict):
        raise ValueError("必须提供有效的正则表达式模式字典")
    
    # 编译正则表达式
    compiled_patterns = {category: re.compile(pattern) for category, pattern in patterns.items()}
    
    # 用于存储每个类别移动的文件数
    stats = {category: 0 for category in patterns}
    stats["未匹配"] = 0
    
    # 查找源目录中的所有文件（不包括子目录中的文件）
    files = [f for f in source_path.iterdir() if f.is_file()]
    
    # 按正则表达式整理
    for file_path in files:
        file_name = file_path.name
        target_category = "未匹配"
        
        # 检查文件名是否匹配任何模式
        for category, pattern in compiled_patterns.items():
            if pattern.search(file_name):
                target_category = category
                break
        
        # 创建目标子目录（如果不存在）
        target_subdir = target_dir / target_category
        if not target_subdir.exists():
            target_subdir.mkdir(parents=True)
        
        # 移动文件
        try:
            target_file = target_subdir / file_path.name
            if target_file.exists():
                # 处理重名文件
                name_parts = file_path.stem, file_path.suffix
                target_file = target_subdir / f"{name_parts[0]}_copy{name_parts[1]}"
                
            shutil.move(str(file_path), str(target_file))
            stats[target_category] += 1
        except Exception as e:
            print(f"移动文件 {file_path} 时出错: {e}")
    
    return stats 